#!/usr/bin/python
import numpy
import pylab
import util
import os
import cmath
from scipy import signal

def getz(l, i):
    try:
        return l[i]
    except:
        return 0 

def mylfilter(b, a, data):
    """
    a[0]y[n]+a[1]y(n-1)+ ... + a[N]y(n-N) = b[0]x(n) + b[1]x(n-1) + ... + b[n]x(n-N)
    """
    new_data = [0.0]*(len(data))
    for n in xrange(len(data)):
        for k in xrange(len(b)):
            new_data[n] += b[k]*getz(data, n-k)
        for k in xrange(1, len(a)):
            new_data[n] -= a[k]*getz(new_data, n-k) 
        new_data[n] = new_data[n]/a[0]
    return new_data

def Ryy(iq_data, n, d, m):
    """
        iq_data: the date to be calculated
        n: start index
        m: integration length
        d: correlation lag
    """
    if d==0:
        return numpy.array(
            [   numpy.sum(
                    numpy.abs(iq_data[i+n:i+n+m])**2 
                )
                
              for i in xrange(1000)
            ]
        )
    else:
        return numpy.array(
            [ numpy.sum(iq_data[i+n:i+n+m] * numpy.conjugate(iq_data[i+n+d:i+n+d+m]))
              for i in xrange(1000)
            ]
        )
        

def Xyy(d0, d1, n, d, m):
    return numpy.array(
        [   numpy.sum(
                d0[i+n:i+n+m] * numpy.conjugate(d1[i+n+d:i+n+d+m])
            )
            for i in xrange(1000)
        ]
    )

def x(data):
    return range(len(data))

def denominator(iq_data, n, D, M):
    return (Ryy(iq_data, 0, 0, M)+Ryy(iq_data, D, 0, M))**2
    
def dx(d0, d1, n, d, m):    
    return (Xyy(d0, d0, 0, 0, m) + Xyy(d1, d1, d, 0, m))**2

def numerator(iq_data, n, D, M):
    return  (numpy.abs(Ryy(iq_data, 0, D, M))**2)*2

def nx(d0, d1, n, d, m):
    return (numpy.abs(Xyy(d0, d1, 0, d, m))**2)*2

class Ms():

    def __init__(self, fftsize, cp, debug=False):
        self.reg = []
        self.userdata = []
        self.fft_size = fftsize
        self.cp_size = cp
        self.debug = debug

    def max_gain(self, iq_data):
        tmp = []
        for data in iq_data:
            if numpy.real(data)==0:
                re=0
            elif numpy.imag(data)==0:
                im=0
            else:
                if numpy.real(data) > 0:
                    re = 1
                else:
                    re = -1
                if numpy.imag(data) > 0:
                    im = 1
                else:
                    im = -1
            tmp.append(re+im*cmath.sqrt(-1))
        return tmp

    #delay can be 0,1,2
    def search_preamble(self, iq_data, D, M, delay):
        if len(iq_data) < 1024:
            return -1
        #D correlation lag
        #M length of sigma

        #shifted = signal.signaltools.lfilter([3-delay,delay], [3], iq_data) #1/3 delay

        Mmm  = numerator(iq_data, 0, 341,  341 )/denominator(iq_data, 0, 341,  341) 
        Mmm0  = numerator(iq_data, 0, 683,  341 )/denominator(iq_data, 0, 683,  341) 
        shifted = signal.signaltools.lfilter([2, 1], [3], iq_data) 
        Mmms  = nx(iq_data, shifted, 0, 341,  341 )/dx(iq_data, shifted, 0, 341,  341) 
        shifted = signal.signaltools.lfilter([1, 2], [3], iq_data)
        Mmms1  = nx(iq_data, shifted, 0, 684,  341 )/dx(iq_data, shifted, 0, 684,  341) 
        Mcpn = numerator(iq_data, 0, 1024, 128 )/denominator(iq_data, 0, 1024, 128) 
        Mcpn1 = nx(iq_data, shifted, 0, 128, 128 )/dx(iq_data, shifted, 0, 128, 128) 

        pylab.plot(
            x(Mmm), Mmm, 'b-',
            x(Mmm0), Mmm0, 'b+',
            x(Mmms), Mmms, 'r-',
            x(Mmms1), Mmms1, 'r+',
            x(Mcpn), Mcpn, 'g-',
            x(Mcpn1), Mcpn1, 'g+'
        )
        
        pylab.show()
        os._exit(2)
        X = numpy.correlate(iq_data, numpy.conjugate(iq_data), mode='full')
        pylab.figure()
        pylab.plot(numpy.abs(X)) 

        pylab.figure()
        if max(div) > 100:
            return 1
        else:
            return -1

    def detectpreamble(self, iq_data, prmb_index, prmb_time_offset):
        """
        Msc(n) = (|Ryy(n,L)|**2)/(R2yy(n,0)**2)
        Ryy(n,d)=SIGMA(y(n+m)*conjugate(y(n+m+d)), m=0->(M-1))
        Ryy(n+1, d)=Ryy(n,d) + conjugate(y(n+d))*y(n+d+M)-conjugate(y(n))*y(n+M)
        Ryy(n,0) = (1/M)*(SIGMA((|y(n+m)|)**2), m=0->(M-1))
        Fractional Delay Filter
        3*y(n+a/3)=(3-a)*y(n)+a*y(n+1)
        """
        self.search_preamble(iq_data, 342, 342, 2)
        os._exit(1)
        #while(start_point < len(iq_data)-D-M):
        #    if search_preamble(iq_data[start_point:], 342, 342, 2) > 0:
        #        break
        #    else:
        #        start_point += 1024
            
            
        

        
        if self.debug:pylab.plot(range(len(iq_data)), numpy.real(iq_data), 'r.', range(len(iq_data)), numpy.imag(iq_data), 'b.')
        if self.debug:pylab.figure()
        gama = [0]*self.fft_size
        fi   = [0]*self.fft_size

        for i in range(0, self.fft_size):
            for j in range(self.cp_size): 
                gama[i] =  gama[i] + iq_data[i+j]*numpy.conjugate(iq_data[i+j+self.fft_size])
                fi[i]   = fi[i] + (abs(iq_data[i+j])*abs(iq_data[i+j])+abs(iq_data[i+j+self.fft_size])*abs(iq_data[i+j+self.fft_size]))/2 
        delt = [] 
        for i in range(1024):
            delt.append(gama[i]-fi[i])
        max_value = delt[0]
        index = 0
        for i in range(1024):
            if delt[i] > max_value:
                max_value = delt[i]
                index = i
                
        if index != prmb_time_offset:
            err_msg = '\tprmb_time_offset should be %d, but %d, delt[got] is %x, delt[golden] is %x'%(prmb_time_offset, index, delt[index], delt[prmb_time_offset])
            print err_msg
            pass
        else:
            pass
        max_index = index
                
        if self.debug:pylab.plot(range(len(gama)), gama, 'r.', range(len(fi)), fi, 'g.', range(len(delt)), delt, 'b.')

        preamble_freq = pylab.fft(iq_data[index+128:index+1024+128])
        if self.debug:pylab.figure()
        if self.debug:pylab.plot(range(self.fft_size), numpy.real(preamble_freq), 'bo')
        max_value = -100
        found_index = -1
        result= []
        for i in range(114):
            golden = util.get_preamble(self.fft_size, i, True)
            cor_value = numpy.correlate(preamble_freq, golden)[0]
            result.append(numpy.real(cor_value))
            if abs(cor_value)> max_value:
                max_value = abs(cor_value)
                found_index = i
        if found_index > -1:        
            if prmb_index != found_index:
                print "ERROR prmb index should be %d, but %d"%(prmb_index, found_index)
            else:
                print "found index %d"%found_index
        else:
            print "ERROR, no preamble index found"

        if self.debug:pylab.figure()
        if self.debug:pylab.plot(result)
        if self.debug:pylab.show() 

        
    def rx(self, iq_data, prmb_index, prmb_time_offset):
        self.detectpreamble(iq_data, prmb_index, prmb_time_offset)

